package org.gjy.m8.bucket;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 线程安全的 Token Bucket（支持突发容量）。
 * <p>
 * 4.1 思路
 * 桶容量 capacity，每 refillRate（tokens / sec）向桶中 补充令牌。
 * 请求到来时尝试 抢走 1 张令牌：
 * 若 桶中还有令牌 → 通过。
 * 否则 → 限流。
 *
 * @author gjy
 * @version 1.0
 * @since 2025-10-13 17:57:02
 */
public class TokenBucketRateLimiter<K> implements RateLimiter<K> {

    private static class Bucket {
        final long capacity;               // 最大令牌数（突发上限）
        final double refillRate;           // 每秒补充令牌数
        final AtomicLong tokens;// 当前可用令牌（使用 double 支持小数精度）
        volatile long lastRefillTs;        // 最后一次补充的时间戳（毫秒）

        Bucket(long capacity, double refillRate) {
            this.capacity = capacity;
            this.refillRate = refillRate;
            this.tokens = new AtomicLong(capacity);        // 初始满桶
            this.lastRefillTs = System.nanoTime();
        }

        /**
         * 按当前时间补充令牌（懒惰刷新）
         */
        void refill() {
            long now = System.nanoTime();
            long elapsedNs = now - lastRefillTs;
            long added = (long) ((elapsedNs / 1_000_000_000.0) * refillRate);
            if (added > 0) {
                long token = Math.min(capacity, tokens.get() + added);
                tokens.set(token);
                lastRefillTs = now;
            }
        }

        /**
         * 尝试获取 1 颗令牌（成功返回 true）
         */
        synchronized boolean tryConsume() {
            refill();
            if (tokens.get() >= 1) {
                tokens.decrementAndGet();
                return true;
            }
            return false;
        }
    }

    private final ConcurrentMap<K, Bucket> buckets = new ConcurrentHashMap<>();
    private final long capacity;
    private final double refillRate;   // token / second

    public TokenBucketRateLimiter(long capacity, double refillRate) {
        if (capacity <= 0 || refillRate <= 0) {
            throw new IllegalArgumentException("capacity & refillRate 必须 > 0");
        }
        this.capacity = capacity;
        this.refillRate = refillRate;
    }

    @Override
    public boolean tryAcquire(K key) {
        Bucket b = buckets.computeIfAbsent(key, k -> new Bucket(capacity, refillRate));
        return b.tryConsume();
    }
}
